"""Automated chart generation for data frames."""
import itertools

import IPython


_DATAFRAME_REGISTRY = None


def find_charts(
    df,
    max_chart_instances=None,
):
  """Finds charts compatible with dtypes of the given data frame.

  Args:
    df: (pd.DataFrame) A dataframe.
    max_chart_instances: (int) For a single chart type, the max number instances
      to generate.

  Returns:
    (iterable<ChartSection>) A sequence of chart sections.
  """
  # Lazy import to avoid loading altair and transitive deps on kernel init.
  from google.colab import _quickchart_helpers  # pylint: disable=g-import-not-at-top

  def _ensure_dataframe_registry():
    global _DATAFRAME_REGISTRY
    if _DATAFRAME_REGISTRY is None:
      if IPython.get_ipython():
        variable_namespace = IPython.get_ipython().user_ns
      else:  # Fallback to placeholder namespace in testing environment.
        variable_namespace = {}
      _DATAFRAME_REGISTRY = _quickchart_helpers.DataframeRegistry(
          variable_namespace
      )

  _ensure_dataframe_registry()

  chart_sections = determine_charts(
      df, _DATAFRAME_REGISTRY, max_chart_instances
  )
  if not chart_sections:
    print('No charts were generated by quickchart')
  return chart_sections


def find_charts_json(df_name: str, max_chart_instances=None):
  """Equivalent to find_charts, but emits to JSON for use from browser."""

  class FixedDataframeRegistry:

    def get_or_register_varname(self, _) -> str:
      """Returns the name of the fixed dataframe name."""
      return df_name

  dataframe = IPython.get_ipython().user_ns[df_name]

  chart_sections = determine_charts(
      dataframe, FixedDataframeRegistry(), max_chart_instances
  )
  return IPython.display.JSON([s.to_json() for s in chart_sections])


def determine_charts(df, dataframe_registry, max_chart_instances=None):
  """Finds charts compatible with dtypes of the given data frame."""
  # Lazy import to avoid loading matplotlib and transitive deps on kernel init.
  from google.colab import _quickchart_dtypes  # pylint: disable=g-import-not-at-top
  from google.colab import _quickchart_helpers  # pylint: disable=g-import-not-at-top

  dtype_groups = _quickchart_dtypes.classify_dtypes(df)
  numeric_cols = dtype_groups['numeric']
  categorical_cols = dtype_groups['categorical']
  time_cols = dtype_groups['datetime'] + dtype_groups['timelike']
  chart_sections = []

  if numeric_cols:
    section = _quickchart_helpers.histograms_section(
        df, numeric_cols[:max_chart_instances], dataframe_registry
    )
    if section.charts:
      chart_sections.append(section)

  if categorical_cols:
    selected_categorical_cols = categorical_cols[:max_chart_instances]
    section = _quickchart_helpers.categorical_histograms_section(
        df, selected_categorical_cols, dataframe_registry
    )
    if section.charts:
      chart_sections.append(section)

  if len(numeric_cols) >= 2:
    section = _quickchart_helpers.scatter_section(
        df,
        _select_first_k_pairs(numeric_cols, k=max_chart_instances),
        dataframe_registry,
    )
    if section.charts:
      chart_sections.append(section)

  if time_cols:
    section = _quickchart_helpers.time_series_line_plots_section(
        df,
        _select_time_series_cols(
            time_cols=time_cols,
            numeric_cols=numeric_cols,
            categorical_cols=categorical_cols,
            k=max_chart_instances,
        ),
        dataframe_registry,
    )
    if section.charts:
      chart_sections.append(section)

  if numeric_cols:
    section = _quickchart_helpers.value_plots_section(
        df, numeric_cols[:max_chart_instances], dataframe_registry
    )
    if section.charts:
      chart_sections.append(section)

  if len(categorical_cols) >= 2:
    section = _quickchart_helpers.heatmaps_section(
        df,
        _select_first_k_pairs(categorical_cols, k=max_chart_instances),
        dataframe_registry,
    )
    if section.charts:
      chart_sections.append(section)

  if categorical_cols and numeric_cols:
    section = _quickchart_helpers.faceted_distributions_section(
        df,
        _select_faceted_numeric_cols(
            numeric_cols, categorical_cols, k=max_chart_instances
        ),
        dataframe_registry,
    )
    if section.charts:
      chart_sections.append(section)

  return chart_sections


def _select_first_k_pairs(colnames, k=None):
  """Selects the first k pairs of column names, sequentially.

  e.g., ['a', 'b', 'c'] => [('a', b'), ('b', 'c')] for k=2

  Args:
    colnames: (iterable<str>) Column names from which to generate pairs.
    k: (int) The number of column pairs.

  Returns:
    (list<(str, str)>) A k-length sequence of column name pairs.
  """
  return itertools.islice(itertools.pairwise(colnames), k)


def _select_faceted_numeric_cols(numeric_cols, categorical_cols, k=None):
  """Selects numeric columns and corresponding categorical facets.

  Args:
    numeric_cols: (iterable<str>) Available numeric columns.
    categorical_cols: (iterable<str>) Available categorical columns.
    k: (int) The number of column pairs to select.

  Returns:
    (iter<(str, str)>) Prioritized sequence of (numeric, categorical) column
    pairs.
  """
  return itertools.islice(itertools.product(numeric_cols, categorical_cols), k)


def _select_time_series_cols(time_cols, numeric_cols, categorical_cols, k=None):
  """Selects combinations of colnames that can be plotted as time series.

  Args:
    time_cols: (iter<str>) Available time-like columns.
    numeric_cols: (iter<str>) Available numeric columns.
    categorical_cols: (iter<str>) Available categorical columns.
    k: (int) The number of combinations to select.

  Returns:
    (iter<(str, str, str)>) Prioritized sequence of (time, value, series)
    colname combinations.
  """
  numeric_cols = [c for c in numeric_cols if c not in time_cols]
  numeric_aggregates = ['count()']
  if not categorical_cols:
    categorical_cols = [None]
  return itertools.islice(
      itertools.product(
          time_cols, numeric_cols + numeric_aggregates, categorical_cols
      ),
      k,
  )
